/**
 * Intermediate value模型实现
 * 
 * 实现6种Algorithm的Intermediate valueCalculateFunction
 */

#include "intermediate_models.h"
#include "intermediate_config.h"
#include <string.h>

// ========================================================================
// AES S-box表
// ========================================================================

const uint8_t AES_SBOX[256] = {
    0x63, 0x7c, 0x77, 0x7b, 0xf2, 0x6b, 0x6f, 0xc5, 0x30, 0x01, 0x67, 0x2b, 0xfe, 0xd7, 0xab, 0x76,
    0xca, 0x82, 0xc9, 0x7d, 0xfa, 0x59, 0x47, 0xf0, 0xad, 0xd4, 0xa2, 0xaf, 0x9c, 0xa4, 0x72, 0xc0,
    0xb7, 0xfd, 0x93, 0x26, 0x36, 0x3f, 0xf7, 0xcc, 0x34, 0xa5, 0xe5, 0xf1, 0x71, 0xd8, 0x31, 0x15,
    0x04, 0xc7, 0x23, 0xc3, 0x18, 0x96, 0x05, 0x9a, 0x07, 0x12, 0x80, 0xe2, 0xeb, 0x27, 0xb2, 0x75,
    0x09, 0x83, 0x2c, 0x1a, 0x1b, 0x6e, 0x5a, 0xa0, 0x52, 0x3b, 0xd6, 0xb3, 0x29, 0xe3, 0x2f, 0x84,
    0x53, 0xd1, 0x00, 0xed, 0x20, 0xfc, 0xb1, 0x5b, 0x6a, 0xcb, 0xbe, 0x39, 0x4a, 0x4c, 0x58, 0xcf,
    0xd0, 0xef, 0xaa, 0xfb, 0x43, 0x4d, 0x33, 0x85, 0x45, 0xf9, 0x02, 0x7f, 0x50, 0x3c, 0x9f, 0xa8,
    0x51, 0xa3, 0x40, 0x8f, 0x92, 0x9d, 0x38, 0xf5, 0xbc, 0xb6, 0xda, 0x21, 0x10, 0xff, 0xf3, 0xd2,
    0xcd, 0x0c, 0x13, 0xec, 0x5f, 0x97, 0x44, 0x17, 0xc4, 0xa7, 0x7e, 0x3d, 0x64, 0x5d, 0x19, 0x73,
    0x60, 0x81, 0x4f, 0xdc, 0x22, 0x2a, 0x90, 0x88, 0x46, 0xee, 0xb8, 0x14, 0xde, 0x5e, 0x0b, 0xdb,
    0xe0, 0x32, 0x3a, 0x0a, 0x49, 0x06, 0x24, 0x5c, 0xc2, 0xd3, 0xac, 0x62, 0x91, 0x95, 0xe4, 0x79,
    0xe7, 0xc8, 0x37, 0x6d, 0x8d, 0xd5, 0x4e, 0xa9, 0x6c, 0x56, 0xf4, 0xea, 0x65, 0x7a, 0xae, 0x08,
    0xba, 0x78, 0x25, 0x2e, 0x1c, 0xa6, 0xb4, 0xc6, 0xe8, 0xdd, 0x74, 0x1f, 0x4b, 0xbd, 0x8b, 0x8a,
    0x70, 0x3e, 0xb5, 0x66, 0x48, 0x03, 0xf6, 0x0e, 0x61, 0x35, 0x57, 0xb9, 0x86, 0xc1, 0x1d, 0x9e,
    0xe1, 0xf8, 0x98, 0x11, 0x69, 0xd9, 0x8e, 0x94, 0x9b, 0x1e, 0x87, 0xe9, 0xce, 0x55, 0x28, 0xdf,
    0x8c, 0xa1, 0x89, 0x0d, 0xbf, 0xe6, 0x42, 0x68, 0x41, 0x99, 0x2d, 0x0f, 0xb0, 0x54, 0xbb, 0x16
};

// ========================================================================
// SM4 S-box表
// ========================================================================

const uint8_t SM4_SBOX[256] = {
    0xD6, 0x90, 0xE9, 0xFE, 0xCC, 0xE1, 0x3D, 0xB7, 0x16, 0xB6, 0x14, 0xC2, 0x28, 0xFB, 0x2C, 0x05,
    0x2B, 0x67, 0x9A, 0x76, 0x2A, 0xBE, 0x04, 0xC3, 0xAA, 0x44, 0x13, 0x26, 0x49, 0x86, 0x06, 0x99,
    0x9C, 0x42, 0x50, 0xF4, 0x91, 0xEF, 0x98, 0x7A, 0x33, 0x54, 0x0B, 0x43, 0xED, 0xCF, 0xAC, 0x62,
    0xE4, 0xB3, 0x1C, 0xA9, 0xC9, 0x08, 0xE8, 0x95, 0x80, 0xDF, 0x94, 0xFA, 0x75, 0x8F, 0x3F, 0xA6,
    0x47, 0x07, 0xA7, 0xFC, 0xF3, 0x73, 0x17, 0xBA, 0x83, 0x59, 0x3C, 0x19, 0xE6, 0x85, 0x4F, 0xA8,
    0x68, 0x6B, 0x81, 0xB2, 0x71, 0x64, 0xDA, 0x8B, 0xF8, 0xEB, 0x0F, 0x4B, 0x70, 0x56, 0x9D, 0x35,
    0x1E, 0x24, 0x0E, 0x5E, 0x63, 0x58, 0xD1, 0xA2, 0x25, 0x22, 0x7C, 0x3B, 0x01, 0x21, 0x78, 0x87,
    0xD4, 0x00, 0x46, 0x57, 0x9F, 0xD3, 0x27, 0x52, 0x4C, 0x36, 0x02, 0xE7, 0xA0, 0xC4, 0xC8, 0x9E,
    0xEA, 0xBF, 0x8A, 0xD2, 0x40, 0xC7, 0x38, 0xB5, 0xA3, 0xF7, 0xF2, 0xCE, 0xF9, 0x61, 0x15, 0xA1,
    0xE0, 0xAE, 0x5D, 0xA4, 0x9B, 0x34, 0x1A, 0x55, 0xAD, 0x93, 0x32, 0x30, 0xF5, 0x8C, 0xB1, 0xE3,
    0x1D, 0xF6, 0xE2, 0x2E, 0x82, 0x66, 0xCA, 0x60, 0xC0, 0x29, 0x23, 0xAB, 0x0D, 0x53, 0x4E, 0x6F,
    0xD5, 0xDB, 0x37, 0x45, 0xDE, 0xFD, 0x8E, 0x2F, 0x03, 0xFF, 0x6A, 0x72, 0x6D, 0x6C, 0x5B, 0x51,
    0x8D, 0x1B, 0xAF, 0x92, 0xBB, 0xDD, 0xBC, 0x7F, 0x11, 0xD9, 0x5C, 0x41, 0x1F, 0x10, 0x5A, 0xD8,
    0x0A, 0xC1, 0x31, 0x88, 0xA5, 0xCD, 0x7B, 0xBD, 0x2D, 0x74, 0xD0, 0x12, 0xB8, 0xE5, 0xB4, 0xB0,
    0x89, 0x69, 0x97, 0x4A, 0x0C, 0x96, 0x77, 0x7E, 0x65, 0xB9, 0xF1, 0x09, 0xC5, 0x6E, 0xC6, 0x84,
    0x18, 0xF0, 0x7D, 0xEC, 0x3A, 0xDC, 0x4D, 0x20, 0x79, 0xEE, 0x5F, 0x3E, 0xD7, 0xCB, 0x39, 0x48
};

// ========================================================================
// 全局Key存储
// ========================================================================

static uint8_t g_aes_key[16];
static uint8_t g_sm4_key[16];
static uint8_t g_rsa_key[256];
static uint8_t g_ecc_key[32];
static uint8_t g_sm2_key[32];
static uint8_t g_dilithium_key[32];

// ========================================================================
// 辅助Function实现
// ========================================================================

uint32_t hamming_weight(uint32_t value) {
    uint32_t count = 0;
    while (value) {
        count += value & 1;
        value >>= 1;
    }
    return count;
}

uint32_t hamming_distance(uint32_t a, uint32_t b) {
    return hamming_weight(a ^ b);
}

// ========================================================================
// 通用接口实现
// ========================================================================

void intermediate_model_init(const uint8_t* key) {
    // According toConfiguration的AlgorithmCall对应的InitializeFunction
#ifdef ALGORITHM_AES
    aes_model_init(key);
#elif defined(ALGORITHM_SM4)
    sm4_model_init(key);
#elif defined(ALGORITHM_RSA)
    rsa_model_init(key);
#elif defined(ALGORITHM_ECC)
    ecc_model_init(key);
#elif defined(ALGORITHM_SM2)
    sm2_model_init(key);
#elif defined(ALGORITHM_DILITHIUM)
    dilithium_model_init(key);
#endif
}

uint32_t compute_intermediate_value(const uint8_t* input) {
    // According toConfiguration的模型Call对应的CalculateFunction
#ifdef MODEL_AES_SBOX_OUTPUT
    return aes_sbox_output(input);
#elif defined(MODEL_AES_SBOX_INPUT)
    return aes_sbox_input(input);
#elif defined(MODEL_AES_HW_SBOX_OUTPUT)
    return aes_hw_sbox_output(input);
#elif defined(MODEL_AES_HD_PLAINTEXT_SBOX)
    return aes_hd_plaintext_sbox(input);
#elif defined(MODEL_AES_ROUND_KEY_XOR)
    return aes_round_key_xor(input);
#elif defined(MODEL_AES_HW_ROUND_KEY_XOR)
    return aes_hw_round_key_xor(input);
    
#elif defined(MODEL_SM4_SBOX_OUTPUT)
    return sm4_sbox_output(input);
#elif defined(MODEL_SM4_SBOX_INPUT)
    return sm4_sbox_input(input);
#elif defined(MODEL_SM4_HW_SBOX_OUTPUT)
    return sm4_hw_sbox_output(input);
#elif defined(MODEL_SM4_TAU_TRANSFORM)
    return sm4_tau_transform(input);
#elif defined(MODEL_SM4_ROUND_FUNCTION_INPUT)
    return sm4_round_function_input(input);
    
#elif defined(MODEL_RSA_MODULAR_MULT)
    return rsa_modular_multiplication(input);
#elif defined(MODEL_RSA_SQUARE_OP)
    return rsa_square_operation(input);
#elif defined(MODEL_RSA_HW_RESULT)
    return rsa_hw_result(input);
#elif defined(MODEL_RSA_MODULAR_ADD)
    return rsa_modular_addition(input);
    
#elif defined(MODEL_ECC_POINT_ADD_X)
    return ecc_point_addition_x(input);
#elif defined(MODEL_ECC_POINT_DBL_X)
    return ecc_point_doubling_x(input);
#elif defined(MODEL_ECC_SCALAR_MULT_STEP)
    return ecc_scalar_mult_step(input);
#elif defined(MODEL_ECC_HW_COORDINATE)
    return ecc_hw_coordinate(input);
    
#elif defined(MODEL_SM2_POINT_ADD_X)
    return sm2_point_addition_x(input);
#elif defined(MODEL_SM2_POINT_DBL_X)
    return sm2_point_doubling_x(input);
#elif defined(MODEL_SM2_SCALAR_MULT_STEP)
    return sm2_scalar_mult_step(input);
#elif defined(MODEL_SM2_MODULAR_OP)
    return sm2_modular_operation(input);
    
#elif defined(MODEL_DILITHIUM_NTT_BUTTERFLY)
    return dilithium_ntt_butterfly(input);
#elif defined(MODEL_DILITHIUM_POLY_MULT)
    return dilithium_polynomial_mult(input);
#elif defined(MODEL_DILITHIUM_REJECTION_SAMPLING)
    return dilithium_rejection_sampling(input);
#elif defined(MODEL_DILITHIUM_MODULAR_RED)
    return dilithium_modular_reduction(input);
#else
    return 0;  // 未定义模型
#endif
}

// ========================================================================
// AES模型实现
// ========================================================================

void aes_model_init(const uint8_t* key) {
    memcpy(g_aes_key, key, 16);
}

uint32_t aes_sbox_output(const uint8_t* input) {
    // SBOX[P[0] ⊕ K[0]]
    uint8_t sbox_input = input[0] ^ g_aes_key[0];
    return AES_SBOX[sbox_input];
}

uint32_t aes_sbox_input(const uint8_t* input) {
    // P[0] ⊕ K[0]
    return input[0] ^ g_aes_key[0];
}

uint32_t aes_hw_sbox_output(const uint8_t* input) {
    // HW(SBOX[P[0] ⊕ K[0]])
    uint32_t sbox_out = aes_sbox_output(input);
    return hamming_weight(sbox_out);
}

uint32_t aes_hd_plaintext_sbox(const uint8_t* input) {
    // HD(P[0], SBOX[P[0] ⊕ K[0]])
    uint32_t sbox_out = aes_sbox_output(input);
    return hamming_distance(input[0], sbox_out);
}

uint32_t aes_round_key_xor(const uint8_t* input) {
    // P[0] ⊕ K[0] (与sbox_input相同但概念上YesAddRoundKey操作)
    return input[0] ^ g_aes_key[0];
}

uint32_t aes_hw_round_key_xor(const uint8_t* input) {
    // HW(P[0] ⊕ K[0])
    uint32_t xor_result = aes_round_key_xor(input);
    return hamming_weight(xor_result);
}

// ========================================================================
// SM4模型实现
// ========================================================================

void sm4_model_init(const uint8_t* key) {
    memcpy(g_sm4_key, key, 16);
}

uint32_t sm4_sbox_output(const uint8_t* input) {
    // SBOX[P[0] ⊕ K[0]]
    uint8_t sbox_input = input[0] ^ g_sm4_key[0];
    return SM4_SBOX[sbox_input];
}

uint32_t sm4_sbox_input(const uint8_t* input) {
    // P[0] ⊕ K[0]
    return input[0] ^ g_sm4_key[0];
}

uint32_t sm4_hw_sbox_output(const uint8_t* input) {
    // HW(SBOX[P[0] ⊕ K[0]])
    uint32_t sbox_out = sm4_sbox_output(input);
    return hamming_weight(sbox_out);
}

uint32_t sm4_tau_transform(const uint8_t* input) {
    // Tau变换：4个并行S-box（简化版，只取First个字节）
    uint8_t b0 = SM4_SBOX[input[0] ^ g_sm4_key[0]];
    uint8_t b1 = SM4_SBOX[input[1] ^ g_sm4_key[1]];
    uint8_t b2 = SM4_SBOX[input[2] ^ g_sm4_key[2]];
    uint8_t b3 = SM4_SBOX[input[3] ^ g_sm4_key[3]];
    return (b0 << 24) | (b1 << 16) | (b2 << 8) | b3;
}

uint32_t sm4_round_function_input(const uint8_t* input) {
    // 轮FunctionInput：X1 ⊕ X2 ⊕ X3 ⊕ RK（简化）
    return input[0] ^ input[1] ^ input[2] ^ g_sm4_key[0];
}

// ========================================================================
// RSA模型实现（简化版）
// ========================================================================

void rsa_model_init(const uint8_t* key) {
    memcpy(g_rsa_key, key, 32);  // 只Use前32字节
}

uint32_t rsa_modular_multiplication(const uint8_t* input) {
    // 简化的模乘：(input * key) mod 0xFFFFFFFF
    uint32_t in_val = (input[0] << 24) | (input[1] << 16) | (input[2] << 8) | input[3];
    uint32_t key_val = (g_rsa_key[0] << 24) | (g_rsa_key[1] << 16) | (g_rsa_key[2] << 8) | g_rsa_key[3];
    uint64_t result = ((uint64_t)in_val * key_val);
    return (uint32_t)(result & 0xFFFFFFFF);
}

uint32_t rsa_square_operation(const uint8_t* input) {
    // 平方操作：input^2 mod 0xFFFFFFFF
    uint32_t in_val = (input[0] << 24) | (input[1] << 16) | (input[2] << 8) | input[3];
    uint64_t result = ((uint64_t)in_val * in_val);
    return (uint32_t)(result & 0xFFFFFFFF);
}

uint32_t rsa_hw_result(const uint8_t* input) {
    // 模乘Result的汉明权重
    uint32_t result = rsa_modular_multiplication(input);
    return hamming_weight(result);
}

uint32_t rsa_modular_addition(const uint8_t* input) {
    // 模加：(input + key) mod 0xFFFFFFFF
    uint32_t in_val = (input[0] << 24) | (input[1] << 16) | (input[2] << 8) | input[3];
    uint32_t key_val = (g_rsa_key[0] << 24) | (g_rsa_key[1] << 16) | (g_rsa_key[2] << 8) | g_rsa_key[3];
    return in_val + key_val;  // 32位自然溢出i.e.For模2^32
}

// ========================================================================
// ECC模型实现（简化版）
// ========================================================================

void ecc_model_init(const uint8_t* key) {
    memcpy(g_ecc_key, key, 32);
}

uint32_t ecc_point_addition_x(const uint8_t* input) {
    // 简化的点加法X坐标Calculate
    uint32_t x1 = (input[0] << 24) | (input[1] << 16) | (input[2] << 8) | input[3];
    uint32_t x2 = (g_ecc_key[0] << 24) | (g_ecc_key[1] << 16) | (g_ecc_key[2] << 8) | g_ecc_key[3];
    // 简化：s^2 - x1 - x2 (Not进行True实的模运算)
    uint32_t s = (x1 ^ x2);  // 简化的斜率
    return (s * s) - x1 - x2;
}

uint32_t ecc_point_doubling_x(const uint8_t* input) {
    // 简化的点倍乘X坐标
    uint32_t x = (input[0] << 24) | (input[1] << 16) | (input[2] << 8) | input[3];
    uint32_t s = (3 * x * x);  // 简化的斜率
    return (s * s) - (2 * x);
}

uint32_t ecc_scalar_mult_step(const uint8_t* input) {
    // Scalar乘法的一步：According to位决定加法Or倍乘
    uint8_t bit = g_ecc_key[0] & 1;
    if (bit) {
        return ecc_point_addition_x(input);
    } else {
        return ecc_point_doubling_x(input);
    }
}

uint32_t ecc_hw_coordinate(const uint8_t* input) {
    // 坐标的汉明权重
    uint32_t x = ecc_point_addition_x(input);
    return hamming_weight(x);
}

// ========================================================================
// SM2模型实现（简化版，与ECCClass似）
// ========================================================================

void sm2_model_init(const uint8_t* key) {
    memcpy(g_sm2_key, key, 32);
}

uint32_t sm2_point_addition_x(const uint8_t* input) {
    // 与ECC点加法Class似
    uint32_t x1 = (input[0] << 24) | (input[1] << 16) | (input[2] << 8) | input[3];
    uint32_t x2 = (g_sm2_key[0] << 24) | (g_sm2_key[1] << 16) | (g_sm2_key[2] << 8) | g_sm2_key[3];
    uint32_t s = (x1 ^ x2);
    return (s * s) - x1 - x2;
}

uint32_t sm2_point_doubling_x(const uint8_t* input) {
    uint32_t x = (input[0] << 24) | (input[1] << 16) | (input[2] << 8) | input[3];
    uint32_t s = (3 * x * x);
    return (s * s) - (2 * x);
}

uint32_t sm2_scalar_mult_step(const uint8_t* input) {
    uint8_t bit = g_sm2_key[0] & 1;
    if (bit) {
        return sm2_point_addition_x(input);
    } else {
        return sm2_point_doubling_x(input);
    }
}

uint32_t sm2_modular_operation(const uint8_t* input) {
    // 简化的模运算
    uint32_t in_val = (input[0] << 24) | (input[1] << 16) | (input[2] << 8) | input[3];
    uint32_t key_val = (g_sm2_key[0] << 24) | (g_sm2_key[1] << 16) | (g_sm2_key[2] << 8) | g_sm2_key[3];
    return (in_val * key_val) ^ (in_val + key_val);
}

// ========================================================================
// Dilithium模型实现（简化版）
// ========================================================================

void dilithium_model_init(const uint8_t* key) {
    memcpy(g_dilithium_key, key, 32);
}

uint32_t dilithium_ntt_butterfly(const uint8_t* input) {
    // 简化的NTT蝶形运算
    uint32_t a = (input[0] << 24) | (input[1] << 16) | (input[2] << 8) | input[3];
    uint32_t b = (input[4] << 24) | (input[5] << 16) | (input[6] << 8) | input[7];
    uint32_t omega = (g_dilithium_key[0] << 24) | (g_dilithium_key[1] << 16) | 
                    (g_dilithium_key[2] << 8) | g_dilithium_key[3];
    
    // Butterfly: u = a + w*b, v = a - w*b (mod q)
    uint32_t wb = (b * omega) & 0x007FFFFF;  // mod 8380417 的简化
    return (a + wb) ^ (a - wb);
}

uint32_t dilithium_polynomial_mult(const uint8_t* input) {
    // 简化的多项式乘法
    uint32_t coeff = (input[0] << 24) | (input[1] << 16) | (input[2] << 8) | input[3];
    uint32_t key_coeff = (g_dilithium_key[0] << 24) | (g_dilithium_key[1] << 16) | 
                         (g_dilithium_key[2] << 8) | g_dilithium_key[3];
    return (coeff * key_coeff) & 0x007FFFFF;
}

uint32_t dilithium_rejection_sampling(const uint8_t* input) {
    // 简化的拒绝Sampling
    uint32_t sample = (input[0] << 24) | (input[1] << 16) | (input[2] << 8) | input[3];
    uint32_t threshold = (g_dilithium_key[0] << 24) | (g_dilithium_key[1] << 16) | 
                         (g_dilithium_key[2] << 8) | g_dilithium_key[3];
    return sample < threshold ? sample : (sample % threshold);
}

uint32_t dilithium_modular_reduction(const uint8_t* input) {
    // 简化的模约简：mod 8380417
    uint32_t value = (input[0] << 24) | (input[1] << 16) | (input[2] << 8) | input[3];
    return value & 0x007FFFFF;  // 简化版
}
